import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from python_ai.common.xcommon import sep

pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False


def x_to_key(x, x_v, y):
    return f'{x}={x_v}|y={y}'


def x_naive_bayes(df, X):
    P = {}

    y_counts_series = df['Y'].value_counts()
    y_p_series = y_counts_series / m  # P of each target value
    Yvs = y_counts_series.index  # all target values

    Xs = df.columns[:-1]  # all features' name

    for y_v in Yvs:
        df2 = df[df['Y'] == y_v]
        len_y = y_counts_series[y_v]
        py = y_p_series[y_v]
        P[y_v] = py

        for x in Xs:
            x_count_series = df2[x].value_counts()
            x_p_series = x_count_series / len_y
            x_vs = x_count_series.index

            for x_v in x_vs:
                k = x_to_key(x, x_v, y_v)
                P[k] = x_p_series[x_v]

    print(P)
    result = []
    for y_v in Yvs:
        r = P[y_v]
        for i, x_v in enumerate(X):
            k = x_to_key(Xs[i], x_v, y_v)
            r *= P[k]
        result.append([y_v, r])
    result = np.array(result)
    idx_sorted = result[:, 1].argsort()[-1::-1]
    result = result[idx_sorted]
    return result[0, 0], result


if '__main__' == __name__:
    sep('load')
    df = pd.read_csv('bayes_lihang.txt', header=0)
    m = len(df)
    print(m)
    print(df[:5])

    X = [2, 'S']
    result = x_naive_bayes(df, X)
    print(result)
